#include "Matrix.h"

Matrix::Matrix()
{
    n = m = 0;
    a = nullptr;
}
Matrix::Matrix(const int &_n)
{
    n = m = _n;
    a = new double[n * m];
    memset(a, 0, sizeof(double) * (n * m));
}
Matrix::Matrix(const int &_n, const int &_m)
{
    n = _n;
    m = _m;
    a = new double[n * m];
    memset(a, 0, sizeof(double) * n * m);
}
Matrix::Matrix(const Matrix &A)
{
    n = A.n;
    m = A.m;
    a = new double[n * m];
    for (int i = 0; i < n; i++)
        for (int j = 0; j < m; j++)
            a[i * m + j] = A(i, j);
}

// 返回Matrix;将传入的数组p构造成n*1的矩阵
Matrix::Matrix(const double *p, const int &_n)
{
    n = _n;
    m = 1;
    a = new double[n];
    for (int i = 0; i < n; i++)
        a[i] = p[i];
}
// 返回Matrix;将传入的数组p构造成n*m的矩阵
Matrix::Matrix(const double *p, const int &_n, const int &_m)
{
    n = _n;
    m = _m;
    a = new double[n * m];
    for (int i = 0; i < n * m; i++)
        a[i] = p[i];
}

Matrix::~Matrix() { delete[] a; }

Matrix &Matrix::operator=(const Matrix &rhs)
{
    n = rhs.n;
    m = rhs.m;
    a = new double[n * m];
    for (int i = 0; i < n; i++)
        for (int j = 0; j < m; j++)
            a[i * m + j] = rhs(i, j);
    return (*this);
}

// 返回矩阵第r行，第c列的常值元素
const double Matrix::element(const int &r, const int &c) const
{
    if (r < 0 || r >= n || c < 0 || c >= m)
    {
        std::cerr << "[Error] Matrix:: out of range" << std::endl;
        exit(-1);
    }
    return a[r * m + c];
}

// 返回矩阵第r行，第c列的元素
double &Matrix::element(const int &r, const int &c)
{
    if (r < 0 || r >= n || c < 0 || c >= m)
    {
        std::cerr << "[Error] Matrix:: out of range" << std::endl;
        exit(-1);
    }
    return a[r * m + c];
}

// 返回矩阵的p范数
double Matrix::vecnorm(const double &p) const
{
    double norm = 0;
    for (int i = 0; i < n * m; i++)
        norm += pow(a[i], p);
    return pow(norm, 1.0 / p);
}
const double Matrix::operator()(const int &r, const int &c) const { return element(r, c); }
double &Matrix::operator()(const int &r, const int &c) { return element(r, c); }

// 返回矩阵第r行
RowVector Matrix::getRow(const int &r) const
{
    RowVector res(m);
    for (int i = 0; i < m; i++)
        res(i) = element(r, i);
    return res;
}
// 返回矩阵第c列
ColVector Matrix::getCol(const int &c) const
{
    ColVector res(n);
    for (int i = 0; i < n; i++)
        res(i) = element(i, c);
    return res;
}

Matrix Matrix::operator+(const Matrix &B) const
{
    if (n != B.n || m != B.m)
    {
        std::cerr << "Dimension conflicts!" << std::endl;
        exit(-1);
    }
    Matrix C(n, m);
    for (int i = 0; i < C.n; i++)
        for (int j = 0; j < C.m; j++)
            C(i, j) = element(i, j) + B(i, j);
    return C;
}
// 矩阵取负
Matrix Matrix::operator-() const
{
    Matrix C(n, m);
    for (int i = 0; i < C.n; i++)
        for (int j = 0; j < C.m; j++)
            C(i, j) = -element(i, j);
    return C;
}
Matrix Matrix::operator-(const Matrix &B) const
{
    if (n != B.n || m != B.m)
    {
        std::cerr << "Dimension conflicts!" << std::endl;
        exit(-1);
    }
    Matrix C(n, m);
    for (int i = 0; i < C.n; i++)
        for (int j = 0; j < C.m; j++)
            C(i, j) = element(i, j) + B(i, j);
    return C;
}
Matrix Matrix::operator*(const Matrix &B) const
{
    if (m != B.n)
    {
        std::cerr << "Dimention Conflicts!" << std::endl;
    }
    Matrix C(n, B.m);
    for (int i = 0; i < C.n; i++)
        for (int j = 0; j < C.m; j++)
            for (int k = 0; k < m; k++)
                C(i, j) += element(i, k) * B(k, j);
    return C;
}

Matrix Matrix::operator/(const double &p) const
{
    Matrix C(n, m);
    for (int i = 0; i < C.n; i++)
        for (int j = 0; j < C.m; j++)
            C(i, j) = element(i, j) / p;
    return C;
}

// 矩阵转置
Matrix Matrix::T() const
{
    Matrix C(m, n);
    for (int i = 0; i < n; i++)
        for (int j = 0; j < m; j++)
            C(j, i) = element(i, j);
    return C;
}

// 交换矩阵第r1行与第r2行
void Matrix::swaprow(const int &r1, const int &r2)
{
    if (r1 < 0 || r1 >= n || r2 < 0 || r2 >= n)
    {
        std::cerr << "Swaprow out of range!" << std::endl;
        exit(-1);
    }
    for (int j = 0; j < m; j++)
        std::swap(a[r1 * m + j], a[r2 * m + j]);
}

// 交换矩阵第c1列和第c2列
void Matrix::swapcol(const int &c1, const int &c2)
{
    if (c1 < 0 || c1 >= m || c2 < 0 || c2 >= m)
    {
        std::cerr << "Swapcol out of range!" << std::endl;
        exit(-1);
    }
    for (int i = 0; i < n; i++)
        std::swap(a[i * m + c1], a[i * m + c2]);
}

// 返回矩阵对角阵
Matrix diag(const Matrix &A)
{
    if (A.n == 1 && A.m > 0)
    {
        Matrix D(A.m, A.m);
        for (int i = 0; i < A.m; i++)
            D(i, i) = A(0, i);
        return D;
    }
    else if (A.m == 1 && A.n > 0)
    {
        Matrix D(A.n, A.n);
        for (int i = 0; i < A.n; i++)
            D(i, i) = A(i, 0);
        return D;
    }
    else
    {
        int n = std::min(A.n, A.m);
        Matrix D(n, 1);
        for (int i = 0; i < n; i++)
            D(i, 0) = A(i, i);
        return D;
    }
}

// 返回矩阵的r1-r2行，c1-c2列的分块矩阵
Matrix Matrix::getSubmatrix(const int &r1, const int &r2, const int &c1, const int &c2) const
{
    if (r1 < 0 || c1 < 0 || r2 >= n || c2 >= m || r2 < r1 || c2 < c1)
    {
        std::cerr << "getSubmatrix:: Out of range!" << std::endl;
        exit(-1);
    }
    Matrix sub(r2 - r1 + 1, c2 - c1 + 1);
    for (int i = 0; i < r2 - r1 + 1; i++)
        for (int j = 0; j < c2 - c1 + 1; j++)
            sub(i, j) = element(r1 + i, c1 + j);
    return sub;
}
// 设置以第(r,c)个元素为左上角的子矩阵为rhs
void Matrix::setSubmatrix(const int &r, const int &c, const Matrix &rhs)
{
    if (r < 0 || c < 0 || r + rhs.n > n || c + rhs.m > m)
    {
        std::cerr << "setSubmatrix:: Out of range!" << std::endl;
        exit(-1);
    }
    for (int i = 0; i < rhs.n; i++)
        for (int j = 0; j < rhs.m; j++)
            element(r + i, j + c) = rhs(i, j);
}

// 返回矩阵的行列式，用高斯消元法计算.
double Matrix::det() const
{
    if (m != n || m == 0)
    {
        std::cerr << "Ill-posed Matrix! Not squared matrix." << std::endl;
        return 0;
    }
    Matrix A(*this);
    double ans = 1;
    for (int i = 0; i < n; i++)
    {
        int p = i;
        while (p < n && A(p, i) == 0)
            p++;
        if (p == n)
            return 0;
        if (p != i)
            A.swaprow(i, p);
        ans *= A(i, i);
        for (int j = i + 1; j < n; j++)
        {
            double coef = A(j, i) / A(i, i);
            for (int k = i; k < n; k++)
                A(j, k) -= A(i, k) * coef;
        }
    }
    return ans;
}
// 返回矩阵的逆矩阵，用列主元高斯消元法计算。
Matrix Matrix::inv() const
{
    if (m != n || m == 0)
    {
        std::cerr << "Matrix Error! There's no inverse of a non-square or empty matrix!" << std::endl;
        return Matrix();
    }
    Matrix A(n, 2 * n);
    for (int i = 0; i < n; i++)
    {
        A(i, i + n) = 1;
        for (int j = 0; j < n; j++)
            A(i, j) = element(i, j);
    }
    for (int i = 0; i < n; i++)
    {
        int p = i;
        for (int j = i + 1; j < n; j++)
            if (fabs(A(j, i)) > fabs(A(p, i)))
                p = j;
        if (A(p, i) == 0)
        {
            std::cerr << "Matrix Error! There's no inverse of a singular matrix!" << std::endl;
            return Matrix();
        }
        if (p != i)
            A.swaprow(i, p);
        for (int k = i + 1; k < 2 * n; k++)
            A(i, k) /= A(i, i);
        A(i, i) = 1;
        for (int j = 0; j < n; j++)
        {
            if (i == j)
                continue;
            double coef = A(j, i);
            for (int k = i; k < 2 * n; k++)
                A(j, k) -= A(i, k) * coef;
        }
    }
    Matrix ans(n, n);
    for (int i = 0; i < n; i++)
        for (int j = 0; j < n; j++)
            ans(i, j) = A(i, j + n);
    return ans;
}

/******************************线性方程组计算（数值代数）******************************/
double det(const Matrix &A) { return A.det(); }
Matrix inv(const Matrix &A) { return A.inv(); }
ColVector solveLowerTriangular(const Matrix &A, const ColVector &b)
{
    ColVector x = b;
    int m = A.m;
    int n = A.n;
    for (int i = 0; i < n; i++)
    {
        x(i) /= A(i, i);
        for (int j = i + 1; j < A.n; j++)
        {
            x(j) -= A(j, i) * x(i);
        }
    }
    return x;
}
ColVector solveUpperTriangular(const Matrix &A, const ColVector &b)
{
    ColVector x = b;
    for (int i = A.n - 1; i >= 0; i--)
    {
        x(i) /= A(i, i);
        for (int j = i - 1; j >= 0; j--)
        {
            x(j) -= x(i) * A(j, i);
        }
    }
    return x;
}
// 利用列主元Gauss消元法解Ax=b
ColVector solve(Matrix A, ColVector b)
{
    if (A.n < A.m || A.n != b.n || A.m == 0)
    {
        std::cerr << "Ill-posed matrix!" << std::endl;
    }
    int n = A.n, m = A.m;
    ColVector x(m);
    int *columns = new int[m];
    int col_cnt;
    for (int i = 0; i < m; i++)
    {
        int p = i;
        for (int j = i; j < n; j++)
            if (fabs(A(j, i)) > fabs(A(p, i)))
                p = j;
        if (p != i)
            A.swaprow(i, p), b.swaprow(i, p);
        if (!A(i, i))
        {
            std::cerr << "Matrix Error! The method solve() cannot solve an singular equation!" << std::endl;
            std::cerr << "Error occurs at i=" << i << std::endl;
            return Matrix();
        }
        col_cnt = 0;
        for (int j = i; j < m; j++)
            if (A(i, j))
                columns[col_cnt++] = j;
        // 冗余列优化，显著提升稀疏矩阵运算速度
        for (int j = i + 1; j < n; j++)
        {
            if (!A(j, i))
                continue; // 冗余行优化，显著提升稀疏矩阵运算速度
            double coef = A(j, i) / A(i, i);
            for (int k = 0; k < col_cnt; k++)
                A(j, columns[k]) -= A(i, columns[k]) * coef;
            b(j) -= b(i) * coef;
        }
    }
    for (int i = m - 1; i >= 0; i--)
    {
        x(i) = b(i) / A(i, i);
        for (int j = 0; j < i; j++)
            b(j) -= A(j, i) * x(i);
    }
    delete[] columns;
    return x;
}
// 改进的Cholesky(LDL)方法，要求矩阵有对称正定的性质。D的值存在L的对角线上。
Matrix choleskyImproved(const Matrix &A)
{
    if (A.m != A.n || A.m == 0)
    {
        std::cerr << "Cholesky Decomposition Failure! Symmetry and positive definiteness are required" << std::endl;
        return Matrix();
    }
    int n = A.n;
    Matrix L(n, n);
    for (int j = 0; j < n; j++)
    {
        L(j, j) = A(j, j);
        for (int k = 0; k < j; k++)
            L(j, j) -= L(j, k) * L(k, k) * L(j, k);
        for (int i = j + 1; i < n; i++)
        {
            L(i, j) = A(i, j);
            for (int k = 0; k < j; k++)
                L(i, j) -= L(i, k) * L(k, k) * L(j, k);
            L(i, j) /= L(j, j);
        }
    }
    return L;
}
// 用改进Cholesky方法求解正定对称方程，用法：x = solveByLDL(A,b)
ColVector solveByLDL(const Matrix &A, const Matrix &b)
{
    Matrix L = choleskyImproved(A);
    Matrix D = diag(L);
    L = L - diag(D) + eye(L.n);
    ColVector y = solveLowerTriangular(L, b);
    for (int i = 0; i < y.n; i++)
        y(i) /= D(i, i);
    return solveUpperTriangular(L.T(), y);
}
// 利用共轭梯度法求解方程组Ax = b, A必须为方阵
ColVector CG_solve(const Matrix &A, const ColVector &b, const double err, ColVector x)
{
    if (x.is_empty())
        x = ColVector(b.size());
    ColVector r = A * x - b;
    ColVector p = -r;
    long long step = 0;
    while (r.vecnorm(2) >= err)
    {
        step++;
        
        double alpha = - r.T() * p / (p.T() * A * p);
        x = x + alpha * p;
        double tmp = r.T() * r;
        r = r + alpha * A * p;
        double beta = r.T() * r / tmp;
        p = -r + beta * p;
    }
    std::cout << "Steps: " << step << std::endl;
    return x;
}
ColVector CG_solve(const Matrix &A, const ColVector &b)
{
    return CG_solve(A, b, 1e-6);
}
ColVector CG_solve(const Matrix &A, const ColVector &b, const double err)
{
    return CG_solve(A, b, err, ColVector());
}

/******************************ColVector******************************/
// 返回列向量；实现n*1矩阵转化为n*1的矩阵
ColVector::ColVector(const Matrix &rhs)
{
    if (rhs.m != 1)
    {
        std::cerr << "Can't convert this Matrix to a ColVector" << std::endl;
        exit(-1);
    }
    (*this) = ColVector(rhs.n);
    for (int i = 0; i < n; i++)
        (*this)(i) = element(i, 0);
}

// 返回列向量的第inx个元素
const double ColVector::operator()(const int &inx) const { return element(inx, 0); }
double &ColVector::operator()(const int &inx) { return element(inx, 0); }

ColVector &ColVector::operator+(const ColVector &rhs) const
{
    if (n != rhs.n)
    {
        std::cerr << "ColVector::+ Dimension Conflicts!" << std::endl;
        exit(-1);
    }
    ColVector res(n);
    for (int i = 0; i < n; i++)
        res(i) = (*this)(i) + rhs(i);
    return res;
}
ColVector &ColVector::operator-(const ColVector &rhs) const
{
    if (n != rhs.n)
    {
        std::cerr << "ColVector::- Dimension Conflicts!" << std::endl;
        exit(-1);
    }
    ColVector res(n);
    for (int i = 0; i < n; i++)
        res(i) = (*this)(i)-rhs(i);
    return res;
}
ColVector &ColVector::operator-() const
{
    ColVector res(n);
    for (int i = 0; i < n; i++)
        res(i) = -(*this)(i);
    return res;
}

// 返回行向量，实现列向量转置
RowVector ColVector::T() const
{
    RowVector res(n);
    for (int i = 0; i < n; i++)
        res(i) = (*this)(i);
    return res;
}

/******************************RowVector******************************/
// 返回行向量；实现将1*n的矩阵化为1*n的行向量
RowVector::RowVector(const Matrix &rhs)
{
    if (rhs.n != 1)
    {
        std::cerr << "Can't convert this Matrix to a RowVector" << std::endl;
        exit(-1);
    }
    (*this) = RowVector(rhs.n);
    for (int i = 0; i < n; i++)
        (*this)(i) = element(i, 0);
}

// 返回行向量的第inx个值
const double RowVector::operator()(const int &inx) const { return element(0, inx); }
double &RowVector::operator()(const int &inx) { return element(0, inx); }

RowVector &RowVector::operator+(const RowVector &rhs) const
{
    if (m != rhs.m)
    {
        std::cerr << "RowVector Addition Error!" << std::endl;
        exit(-1);
    }
    RowVector res(m);
    for (int i = 0; i < m; i++)
        res(i) = (*this)(i) + rhs(i);
    return res;
}
RowVector &RowVector::operator-(const RowVector &rhs) const
{
    if (m != rhs.m)
    {
        std::cerr << "RowVector Addition Error!" << std::endl;
        exit(-1);
    }
    RowVector res(m);
    for (int i = 0; i < m; i++)
        res(i) = (*this)(i)-rhs(i);
    return res;
}
RowVector &RowVector::operator-() const
{
    RowVector res(m);
    for (int i = 0; i < m; i++)
        res(i) = -(*this)(i);
    return res;
}

// 返回列向量；实现行向量转置
ColVector RowVector::T() const
{
    ColVector res(n);
    for (int i = 0; i < n; i++)
        res(i) = (*this)(i);
    return res;
}

/******************************相关运算******************************/
Matrix operator*(const double &k, const Matrix &A)
{
    Matrix C(A.n, A.m);
    for (int i = 0; i < C.n; i++)
        for (int j = 0; j < C.m; j++)
            C(i, j) = k * A(i, j);
    return C;
}
ColVector operator*(const Matrix &A, const ColVector &x)
{
    if (A.m != x.n)
    {
        std::cerr << "Matrix*ColVector:: Dimension Conflicts!" << std::endl;
        return ColVector();
    }
    ColVector res(A.n);
    for (int i = 0; i < res.n; i++)
        for (int j = 0; j < A.m; j++)
            res(i) += A(i, j) * x(j);
    return res;
}
RowVector operator*(const RowVector &x, const Matrix &A)
{
    if (x.m != A.n)
    {
        std::cerr << "RowVector*Matrix:: Dimension Conflicts!" << std::endl;
        return RowVector();
    }
    RowVector res(A.m);
    for (int j = 0; j < A.m; j++)
        for (int i = 0; i < A.n; i++)
            res(i) += x(i) * A(i, j);
    return res;
}
ColVector operator*(const double &k, const ColVector &x)
{
    ColVector res(x.n);
    for (int i = 0; i < res.n; i++)
        res(i) = k * x(i);
    return res;
}
RowVector operator*(const double &k, const RowVector &x)
{
    RowVector res(x.n);
    for (int i = 0; i < x.n; i++)
        res(i) = k * x(i);
    return res;
}
double operator*(const RowVector &rv, const ColVector &cv)
{
    if (rv.size() != cv.size())
    {
        std::cerr << "RowVector*ColVector:: Dimension Conflicts!" << std::endl;
        return 0;
    }
    double res = 0;
    for (int i = 0; i < rv.size(); i++)
    {
        res += rv(i) * cv(i);
    }
    return res;
}
std::istream &operator>>(std::istream &in, const Matrix &A)
{
    for (int i = 0; i < A.n; i++)
        for (int j = 0; j < A.m; j++)
            in >> A(i, j);
    return in;
}
std::ostream &operator<<(std::ostream &out, const Matrix &A)
{
    for (int i = 0; i < A.n; i++)
    {
        out << "[ " << A(i, 0);
        for (int j = 1; j < A.m; j++)
            out << ", " << A(i, j);
        out << " ]";
        if (i < A.n - 1)
            out << std::endl;
    }
    return out;
}

/******************************特殊矩阵（向量）******************************/
Matrix zeros(const int &n, const int &m) { return Matrix(n, m); }
Matrix eye(const int &n)
{
    Matrix A(n, n);
    for (int i = 0; i < n; i++)
        A(i, i) = 1.0;
    return A;
}
ColVector zeroCol(const int &n) { return ColVector(n); }
RowVector zeroRow(const int &n) { return RowVector(n); }